import fire
import json
import random
import os

from metagpt.actions import Action, UserRequirement
from metagpt.const import MESSAGE_ROUTE_TO_NONE
from metagpt.logs import logger
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.team import Team

from calc_24_points import calc_all_expressions


# helper functions here

cached_expressions = {}
recent_point_list = None

def get_random_point_list():
    return [
        random.randint(0, 12) + 1,
        random.randint(0, 12) + 1,
        random.randint(0, 12) + 1,
        random.randint(0, 12) + 1
    ]


def get_cached_expressions(point_list: list):
    global cached_expressions
    
    point_list.sort()
    keyy = json.dumps(point_list)
    
    if keyy in cached_expressions:
        return cached_expressions[keyy]
    
    expressions = calc_all_expressions(point_list)
    cached_expressions[keyy] = expressions
    return expressions


async def get_random_suit() -> list:
    suit_list = ['Clubs', 'Hearts', 'Spades', 'Diamonds']
    random.shuffle(suit_list)
    return suit_list


def get_rank_str(point):
    if point == 1:
        rank = "A"
    elif point == 11:
        rank = "J"
    elif point == 12:
        rank = "Q"
    elif point == 13:
        rank = "K"
    else:
        rank = str(point)
    return rank


def get_recent_point_list(context):
    global recent_point_list
    
    card_list = None
    for item in context:
        if item.role == "user" and '"point_list":' in item.content:
            card_list = item.content
    if card_list is not None:
        recent_point_list = card_list
    elif recent_point_list is not None:
        card_list = recent_point_list
    else:
        print("Unknown error!!!")
        os._exit(1)
        
    point_list = json.loads(card_list)["point_list"]
    return point_list


async def get_random_card_list(number_list):
    suit_list = await get_random_suit()
    
    cards_list = []
    for i in range(4):
        
        point = number_list[i]
        rank = get_rank_str(point)
            
        cards_list.append({"suit": suit_list[i], "rank": rank, "point": point})

    return cards_list


# all action class here


class DealCards(Action):
    name: str = "DealCards"

    async def run(self, context: str):
        point_list = get_random_point_list()
        return f"{{\"point_list\": {point_list}}}"


class GetHumanReply(Action):
    PROMPT_TEMPLATE: str = """
    Cards the dealer just posted: {content}
    Please give an expression for the four operations that results in 24.
    Type 'help' if you feel it's difficult.
    Type 'deal' if you want the dealer to deal cards again.
    Type 'exit' if you want to exit this game, type 'exit'.
    """
    name: str = "GetHumanReply"

    async def run(self, content: str):
        prompt = self.PROMPT_TEMPLATE.format(content=content)
        human_reply = await self._aask(prompt)
        return human_reply


class HumanGiveExpression(Action):
    name: str = "HumanGiveExpression"

    async def run(self, content: str):
        return content


class MachineGiveExpression(Action):
    name: str = "MachineGiveExpression"

    async def run(self, context: str):
        point_list = json.loads(recent_point_list)["point_list"]
        
        if len(point_list) == 0:
            return "expression not found"

        expressions = get_cached_expressions(point_list)
        
        result = "expression not found"
        if len(expressions) > 0:
            random_idx = random.randint(0, len(expressions)-1)
            result = expressions[random_idx]
        
        return result


class CheckExpression(Action):
    name: str = "CheckExpression"

    async def run(self, expression: str):
        if expression == "expression not found":
            return "Correct"
        
        result = eval(expression)
        if abs(result - 24) < 0.001:
            return "Correct"
        else:
            return "Wrong"


class CallMathProdigy(Action):
    name: str = "CallMathProdigy"

    async def run(self, content: str):
        return "CallMathProdigy"


class WrongExpression(Action):
    name: str = "WrongExpression"

    async def run(self, context: str):
        return "WrongExpression"


class RequireDealCardsAgain(Action):
    name: str = "RequireDealCardsAgain"

    async def run(self, content: str):
        return "RequireDealCardsAgain"


class ExitGame(Action):
    name: str = "ExitGame"

    async def run(self, content: str):
        return "ExitGame"


# all role class here

class MathProdigy(Role):
    name: str = "Gauss"
    profile: str = "MathProdigy"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([CallMathProdigy])
        self.set_actions([MachineGiveExpression])


class GameJudger(Role):
    name: str = "Peter"
    profile: str = "GameJudger"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([DealCards])
        self._watch([UserRequirement, RequireDealCardsAgain, HumanGiveExpression, MachineGiveExpression])

    async def _act(self) -> Message:
        context = self.get_memories()
        msg = self.get_memories(k=1)[0]
        cause_by_str = f"{msg.cause_by}"

        if ("UserRequirement" in cause_by_str) or ("RequireDealCardsAgain" in cause_by_str):
            await self.call_deal_cards(context)
        else:
            todo = CheckExpression()
            check_result = await todo.run(msg.content)
            print(f"==== Your expression '{msg.content}' is {check_result}. ====")

            if check_result == "Correct":
                await self.call_deal_cards(context)
            else:
                self.rc.env.publish_message(Message(content=check_result, cause_by=WrongExpression))
            
        return Message(content="dummy message", send_to=MESSAGE_ROUTE_TO_NONE)

    async def call_deal_cards(self, context):
        todo = DealCards()
        card_content = await todo.run(context)
        self.rc.env.publish_message(Message(content=card_content, cause_by=DealCards))


class GamePlayer(Role):
    name: str = "David"
    profile: str = "GamePlayer"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([GetHumanReply])
        self._watch([DealCards, WrongExpression])

    async def _act(self) -> Message:        
        context = self.get_memories() 
        point_list = get_recent_point_list(context)
        card_list = await get_random_card_list(point_list)
        
        todo = self.rc.todo

        cards_content = f"{{'card_list': {card_list}, 'point_list': {point_list}}}"
        human_reply = await todo.run(cards_content)
        
        if human_reply == "deal":
            self.rc.env.publish_message(Message(content="RequireDealCardsAgain", cause_by=RequireDealCardsAgain))
        elif human_reply == "exit":
            self.rc.env.publish_message(Message(content="ExitGame", cause_by=ExitGame))
        elif human_reply == "help":
            self.rc.env.publish_message(Message(content="CallMathProdigy", cause_by=CallMathProdigy))
        else:
            self.rc.env.publish_message(Message(content=human_reply, cause_by=HumanGiveExpression))
            
        return Message(content="dummy message", send_to=MESSAGE_ROUTE_TO_NONE)


async def main(
    idea: str = "play 24 points game with me",
    investment: float = 3.0,
    n_round: int = 500):
    logger.info(idea)

    team = Team(use_mgx=False)
    team.hire(
        [
            MathProdigy(),
            GameJudger(),
            GamePlayer(is_human=True),
        ]
    )

    team.invest(investment=investment)
    team.run_project(idea)
    await team.run(n_round=n_round)


if __name__ == "__main__":
    fire.Fire(main)

